import os, sys
import pickle

wordvec_file_template = 'data/wiki.{}.vec'
dict_file_template = 'data/crosslingual/dictionaries/{}-{}{}.txt'

def load_freq_rank_list(lang):
    wordvec_file = wordvec_file_template.format(lang)
    freq_rank = {}
    try:
        with open(f'dumped/word_freqrank_{lang}.pth', 'rb') as f:
            freq_rank = pickle.load(f)
        return freq_rank
    except Exception:
        pass


    with open(wordvec_file, 'r') as f:
        n, dim  = next(f).split()
        for i, line in enumerate(f):
            word = line.split(' ')[0]
            freq_rank[word] = i+1
            if  (i + 1) % 100000 == 0:
                print("Scanned", i)
            if i >= 200000:
                break

    with open(f'dumped/word_freqrank_{lang}.pth', 'wb') as f:
        pickle.dump(freq_rank, f)


if __name__ == '__main__':
    src = sys.argv[1]
    trg = sys.argv[2]
    train_valid = sys.argv[3]

    if train_valid == 'train':
        dict_file = dict_file_template.format(src, trg, '.0-5000')
    elif train_valid == 'valid':
        dict_file = dict_file_template.format(src, trg, '.5000-6500')
    else:
        dict_file = dict_file_template.format(src, trg, '')
    
    src_rank = load_freq_rank_list(src)
    trg_rank = load_freq_rank_list(trg)
    
    src_dictrk = set()
    trg_dictrk = set()

    with open(dict_file, 'r') as f:
        for line in f:
            srcw , trgw = line.split()
            if srcw == trgw:
                continue
            if srcw in src_rank:
                src_dictrk.add(src_rank[srcw])
            
            if trgw in trg_rank:
                trg_dictrk.add(trg_rank[trgw])
    
    print(src_dictrk)
    print(trg_dictrk)
